{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "### TAO remote client - Classification\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "![image](https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png)\n",
    "\n",
    "### Sample prediction for an Image Classification model\n",
    "<img align=\"center\" src=\"../example_images/sample_image_classification.jpg\">\n",
    "\n",
    "### The workflow in a nutshell\n",
    "\n",
    "- Creating a dataset\n",
    "- Upload dataset to the service\n",
    "- Getting a PTM from NGC\n",
    "- Model Actions\n",
    "    - Train (Normal/AutoML)\n",
    "    - Evaluate\n",
    "    - Prune, retrain\n",
    "    - Export\n",
    "    - Tao-Deploy\n",
    "    - Inference on TAO\n",
    "\n",
    "### Table of contents\n",
    "\n",
    "1. [Install TAO remote client ](#head-1)\n",
    "1. [Set the remote service base URL](#head-2)\n",
    "1. [Access the shared volume](#head-3)\n",
    "1. [Create the datasets](#head-4)\n",
    "1. [List datasets](#head-5)\n",
    "1. [Create an experiment](#head-8)\n",
    "1. [Find pretrained model](#head-9)\n",
    "1. [Customize model metadata](#head-10)\n",
    "1. [View hyperparameters that are enabled for AutoML by default](#head-11)\n",
    "1. [Set AutoML related configurations](#head-12)\n",
    "1. [Provide train specs](#head-13)\n",
    "1. [Run train](#head-14)\n",
    "1. [View checkpoint files](#head-15)\n",
    "1. [Provide evaluate specs](#head-16)\n",
    "1. [Run evaluate](#head-17)\n",
    "1. [Provide prune specs](#head-18)\n",
    "1. [Run prune](#head-19)\n",
    "1. [Provide retrain specs](#head-20)\n",
    "1. [Run retrain](#head-21)\n",
    "1. [Run evaluate on retrain](#head-21-1)\n",
    "1. [Provide export specs](#head-22)\n",
    "1. [Run export](#head-23)\n",
    "1. [Provide trt engine generation specs](#head-26)\n",
    "1. [Run TRT Engine generation using TAO-Deploy](#head-27)\n",
    "1. [Provide TAO inference specs](#head-28)\n",
    "1. [Run TAO inference](#head-29)\n",
    "1. [Delete experiment](#head-30)\n",
    "1. [Delete datasets](#head-31)\n",
    "1. [Unmount shared volume](#head-32)\n",
    "1. [Uninstall TAO Remote Client](#head-33)\n",
    "\n",
    "### Requirements\n",
    "Please find the server requirements [here](https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_api/api_setup.html#)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import subprocess\n",
    "import json\n",
    "import time\n",
    "import ast\n",
    "from IPython.display import clear_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "namespace = 'default'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FIXME\n",
    "\n",
    "1. Assign a model_name in FIXME 1\n",
    "1. Assign a workdir in FIXME 2\n",
    "1. Assign the ip_address and port_number in FIXME 3 ([info](https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_api/api_rest_api.html))\n",
    "1. Assign the ngc_api_key variable in FIXME 4\n",
    "1. (Optional) Enable AutoML if needed in FIXME 5\n",
    "1. (Optional) Choose between bayesian and hyperband automl_algorithm in FIXME 6 (If automl was enabled in FIXME5)\n",
    "1. Choose to download jobs or not in FIXME 7\n",
    "1. Choose between default and custom dataset in FIXME 8\n",
    "1. Assign path of DATA_DIR in FIXME 9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Available models (#FIXME 1):\n",
    "# 1. classification_pyt - https://docs.nvidia.com/tao/tao-toolkit/text/image_classification.html\n",
    "# 2. classification_tf1 - https://docs.nvidia.com/tao/tao-toolkit/text/image_classification.html\n",
    "# 3. classification_tf2 - https://docs.nvidia.com/tao/tao-toolkit/text/image_classification_tf2.html\n",
    "# 4. multitask_classification - https://docs.nvidia.com/tao/tao-toolkit/text/multitask_image_classification.html\n",
    "# classification is the same as multi_class classification\n",
    "\n",
    "model_name = \"multitask_classification\"  # FIXME1 (Add the model name from the above mentioned list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "### Install TAO remote client <a class=\"anchor\" id=\"head-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO-Client wheel.\n",
    "! pip3 install nvidia-transfer-learning-client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# View the version of the TAO-Client\n",
    "! nvtl --version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "### Set the remote service base URL <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define the node_addr and port number\n",
    "workdir = \"workdir_classification\" # FIXME2\n",
    "host_url = \"http://<ip_address>:<port_number>\" # FIXME3 example: https://10.137.149.22:32334\n",
    "# In host machine, node ip_address and port number can be obtained as follows,\n",
    "# ip_address: hostname -i\n",
    "# port_number: kubectl get service ingress-nginx-controller -o jsonpath='{.spec.ports[0].nodePort}'\n",
    "\n",
    "ngc_api_key = \"<ngc_api_key>\" # FIXME4 example: (Add NGC API key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "automl_enabled = False # FIXME5 set to True if you want to run automl for the model chosen in the previous cell\n",
    "automl_algorithm=\"bayesian\" # FIXME6 example: bayesian/hyperband\n",
    "# FIXME7 Defaulted to False as downloading jobs from service to your machine takes time\n",
    "# Set to True if you want to download jobs where examples have been provided like for train, export, inference.\n",
    "download_jobs = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "%env BASE_URL={host_url}/{namespace}/api/v1\n",
    "\n",
    "# Exchange NGC_API_KEY for JWT\n",
    "identity = json.loads(subprocess.getoutput(f\"nvtl login --ngc-api-key {ngc_api_key}\"))\n",
    "\n",
    "%env USER={identity['user_id']}\n",
    "%env TOKEN={identity['token']}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Creating workdir\n",
    "workdir = os.path.abspath(workdir)\n",
    "if not os.path.isdir(workdir):\n",
    "    os.makedirs(workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function to parse logs <a class=\"anchor\" id=\"head-1.1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def my_tail(model_name_cli, id, job_id, job_type, workdir):\n",
    "\tstatus = None\n",
    "\twhile True:\n",
    "\t\ttime.sleep(10)\n",
    "\t\tclear_output(wait=True)\n",
    "\t\tlog_file_path = subprocess.getoutput(f\"nvtl {model_name_cli} get-log-file --id {id} --job {job_id} --job_type {job_type} --workdir {workdir}\")\n",
    "\t\tif not os.path.exists(log_file_path):\n",
    "\t\t\tcontinue\n",
    "\t\twith open(log_file_path, 'rb') as log_file:\n",
    "\t\t\tlog_contents = log_file.read()\n",
    "\t\tlog_content_lines = log_contents.decode(\"utf-8\").split(\"\\n\")        \n",
    "\t\tfor line in log_content_lines:\n",
    "\t\t\tprint(line.strip())\n",
    "\t\t\tif line.strip() == \"Error EOF\":\n",
    "\t\t\t\tstatus = \"Error\"\n",
    "\t\t\t\tbreak\n",
    "\t\t\telif line.strip() == \"Done EOF\":\n",
    "\t\t\t\tstatus = \"Done\"\n",
    "\t\t\t\tbreak\n",
    "\t\tif status is not None:\n",
    "\t\t\tbreak\n",
    "\treturn status"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function to split tar files <a class=\"anchor\" id=\"head-1.1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tarfile\n",
    "\n",
    "def split_tar_file(input_tar_path, output_dir, max_split_size=0.2*1024*1024*1024):\n",
    "\tos.makedirs(output_dir, exist_ok=True)\n",
    "\t\n",
    "\twith tarfile.open(input_tar_path, 'r') as original_tar:\n",
    "\t\tmembers = original_tar.getmembers()\n",
    "\t\tcurrent_split_size = 0\n",
    "\t\tcurrent_split_number = 0\n",
    "\t\tcurrent_split_name = os.path.join(output_dir, f'smaller_file_{current_split_number}.tar')\n",
    "\t\t\n",
    "\t\twith tarfile.open(current_split_name, 'w') as split_tar:\n",
    "\t\t\tfor member in members:\n",
    "\t\t\t\tif current_split_size + member.size <= max_split_size:\n",
    "\t\t\t\t\tsplit_tar.addfile(member, original_tar.extractfile(member))\n",
    "\t\t\t\t\tcurrent_split_size += member.size\n",
    "\t\t\t\telse:\n",
    "\t\t\t\t\tsplit_tar.close()\n",
    "\t\t\t\t\tcurrent_split_number += 1\n",
    "\t\t\t\t\tcurrent_split_name = os.path.join(output_dir, f'smaller_file_{current_split_number}.tar')\n",
    "\t\t\t\t\tcurrent_split_size = 0\n",
    "\t\t\t\t\tsplit_tar = tarfile.open(current_split_name, 'w')  # Open a new split tar archive\n",
    "\t\t\t\t\tsplit_tar.addfile(member, original_tar.extractfile(member))\n",
    "\t\t\t\t\tcurrent_split_size += member.size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set dataset type, format <a class=\"anchor\" id=\"head-1.1\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "**For multi-class classification:**\n",
    "\n",
    "We will be using the `pascal VOC dataset` for the tutorial. To find more details please visit [here](http://host.robots.ox.ac.uk/pascal/VOC/voc2012/index.html#devkit). Please download the dataset present [here](http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar) to the environment variable $DATA_DIR."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**If using custom dataset; it should follow this dataset structure, and skip running** \"**Split dataset into train and val sets**\"\n",
    "```\n",
    "DATA_DIR\n",
    "├── classes.txt\n",
    "├── images_test\n",
    "│   ├── class_name_1\n",
    "│   │   ├── image_name_1.jpg\n",
    "│   │   ├── image_name_2.jpg\n",
    "│   │   ├── ...\n",
    "|   |   ... \n",
    "│   └── class_name_n\n",
    "│       ├── image_name_3.jpg\n",
    "│       ├── image_name_4.jpg\n",
    "│       ├── ...\n",
    "├── images_train\n",
    "│   ├── class_name_1\n",
    "│   │   ├── image_name_5.jpg\n",
    "│   │   ├── image_name_6.jpg\n",
    "|   |   ...\n",
    "│   └── class_name_n\n",
    "│       ├── image_name_7.jpg\n",
    "│       ├── image_name_8.jpg\n",
    "│       ├── ...\n",
    "|\n",
    "└── images_val\n",
    "    ├── class_name_1\n",
    "    │   ├── image_name_9.jpg\n",
    "    │   ├── image_name_10.jpg\n",
    "    │   ├── ...\n",
    "    |   ...\n",
    "    └── class_name_n\n",
    "        ├── image_name_11.jpg\n",
    "        ├── image_name_12.jpg\n",
    "        ├── ...\n",
    "```\n",
    "- Each class name folder should contain the images corresponding to that class\n",
    "- Same class name folders should be present across images_test, images_train and images_val\n",
    "- classes.txt is a file which contains the names of all classes (each name in a separate line)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**For multi-task classification:**\n",
    "\n",
    "We will be using the Fashion Product Images (Small) for the tutorial. This dataset is available on Kaggle.In this tutorial, our trained classification network will perform three tasks: article category classification, base color classification and target season classification.\n",
    "\n",
    "To download the dataset, you will need a Kaggle account. After login, you can download the dataset zip file [here](https://www.kaggle.com/paramaggarwal/fashion-product-images-small). The downloaded file is archive.zip with a subfolder called myntradataset. Unzip contents in this subfolder to your workdir created in the cell above and you should have a folder called images and a CSV file called styles.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**If using custom dataset; it should follow this dataset structure**\n",
    "```\n",
    "DATA_DIR\n",
    "├── images\n",
    "│   ├── image_name_1.jpg\n",
    "│   ├── image_name_2.jpg\n",
    "|   |   ├── ...\n",
    "├── styles.csv\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name == \"classification_pyt\":\n",
    "    ds_format = \"classification_pyt\"\n",
    "elif \"classification_\" in model_name:\n",
    "    ds_format = \"default\"\n",
    "elif model_name == \"multitask_classification\":\n",
    "    ds_format = \"custom\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "dataset_to_be_used = \"default\" # FIXME8 example: default/custom; default for the dataset used in this tutorial notebook; custom for a different dataset\n",
    "DATA_DIR = model_name # FIXME9\n",
    "os.environ['DATA_DIR']= DATA_DIR\n",
    "!mkdir -p $DATA_DIR\n",
    "job_map = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset download and pre-processing <a class=\"anchor\" id=\"head-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if dataset_to_be_used == \"default\":\n",
    "    if \"classification_\" in model_name:\n",
    "        assert os.path.exists(os.path.join(DATA_DIR,\"VOCtrainval_11-May-2012.tar\"))\n",
    "        !tar -xf $DATA_DIR/VOCtrainval_11-May-2012.tar -C $DATA_DIR\n",
    "        assert (os.path.exists(f\"{DATA_DIR}/VOCdevkit/\"))\n",
    "        !rm -rf $DATA_DIR/split\n",
    "    elif model_name == \"multitask_classification\":\n",
    "        assert os.path.exists(os.path.join(DATA_DIR,\"archive.zip\"))\n",
    "        !unzip -uq $DATA_DIR/archive.zip -d $DATA_DIR/\n",
    "        assert (os.path.exists(f\"{DATA_DIR}/images\"))\n",
    "        assert (os.path.exists(f\"{DATA_DIR}/styles.csv\"))\n",
    "        # Create subdirectories and remove existing files in them\n",
    "        !mkdir -p $DATA_DIR/images_train && rm -rf $DATA_DIR/images_train/*\n",
    "        !mkdir -p $DATA_DIR/images_val && rm -rf $DATA_DIR/images_val/*\n",
    "        !mkdir -p $DATA_DIR/images_test && rm -rf $DATA_DIR/images_test/*            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Split dataset into train and val sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!python3 -m pip install numpy pandas==1.5.1 tqdm\n",
    "if \"classification_\" in model_name and dataset_to_be_used == \"default\":\n",
    "    !python3 ../dataset_prepare/classification/dataset_split.py\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/split/images_train/\"))\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/split/images_val/\"))\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/split/images_test/\"))\n",
    "elif model_name == \"multitask_classification\" and dataset_to_be_used == \"default\":\n",
    "    !python3 ../dataset_prepare/multitask_classification/dataset_split.py\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/images_train/\"))\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/images_val/\"))\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/images_test/\"))\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/train.csv\"))\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/val.csv\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Tar files to upload"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "if \"classification_\" in model_name:\n",
    "    !tar -C $DATA_DIR/split/ -czf classification_train.tar.gz images_train classes.txt\n",
    "    !tar -C $DATA_DIR/split/ -czf classification_val.tar.gz images_val classes.txt\n",
    "    !tar -C $DATA_DIR/split/ -czf classification_test.tar.gz images_test classes.txt\n",
    "elif model_name == \"multitask_classification\":\n",
    "    !tar -C $DATA_DIR/ -czf mt_classification_train.tar.gz images_train train.csv val.csv\n",
    "    !tar -C $DATA_DIR/ -czf mt_classification_val.tar.gz images_val val.csv\n",
    "    !tar -C $DATA_DIR/ -czf mt_classification_test.tar.gz images_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "if \"classification_\" in model_name:\n",
    "    train_dataset_path =  \"classification_train.tar.gz\"\n",
    "    eval_dataset_path = \"classification_val.tar.gz\"\n",
    "    test_dataset_path = \"classification_test.tar.gz\"\n",
    "elif model_name == \"multitask_classification\":\n",
    "    train_dataset_path =  \"mt_classification_train.tar.gz\"\n",
    "    eval_dataset_path = \"mt_classification_val.tar.gz\"\n",
    "    test_dataset_path = \"mt_classification_test.tar.gz\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create and upload train dataset <a class=\"anchor\" id=\"head-1.2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_dataset_id = subprocess.getoutput(f\"nvtl {model_name} dataset-create --dataset_type image_classification --dataset_format {ds_format}\")\n",
    "print(train_dataset_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "output_dir = os.path.join(os.path.dirname(os.path.abspath(train_dataset_path)), model_name, \"train\")\n",
    "split_tar_file(train_dataset_path, output_dir)\n",
    "for idx, tar_dataset_path in enumerate(os.listdir(output_dir)):\n",
    "    print(f\"Uploading {idx+1}/{len(os.listdir(output_dir))} tar split\")\n",
    "    upload_train_dataset_message = subprocess.getoutput(f\"nvtl {model_name} dataset-upload --id {train_dataset_id} --path {os.path.join(output_dir,tar_dataset_path)}\")\n",
    "    print(upload_train_dataset_message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create and upload val dataset <a class=\"anchor\" id=\"head-1.3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "eval_dataset_id = subprocess.getoutput(f\"nvtl {model_name} dataset-create --dataset_type image_classification --dataset_format {ds_format}\")\n",
    "print(eval_dataset_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "output_dir = os.path.join(os.path.dirname(os.path.abspath(eval_dataset_path)), model_name, \"val\")\n",
    "split_tar_file(eval_dataset_path, output_dir)\n",
    "for idx, tar_dataset_path in enumerate(os.listdir(output_dir)):\n",
    "    print(f\"Uploading {idx+1}/{len(os.listdir(output_dir))} tar split\")\n",
    "    upload_val_dataset_message = subprocess.getoutput(f\"nvtl {model_name} dataset-upload --id {eval_dataset_id} --path {os.path.join(output_dir,tar_dataset_path)}\")\n",
    "    print(upload_val_dataset_message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create and upload test dataset <a class=\"anchor\" id=\"head-1.4\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "test_dataset_id = subprocess.getoutput(f\"nvtl {model_name} dataset-create --dataset_type image_classification --dataset_format raw\")\n",
    "print(test_dataset_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "output_dir = os.path.join(os.path.dirname(os.path.abspath(test_dataset_path)), model_name, \"test\")\n",
    "split_tar_file(test_dataset_path, output_dir)\n",
    "for idx, tar_dataset_path in enumerate(os.listdir(output_dir)):\n",
    "    print(f\"Uploading {idx+1}/{len(os.listdir(output_dir))} tar split\")\n",
    "    upload_test_dataset_message = subprocess.getoutput(f\"nvtl {model_name} dataset-upload --id {test_dataset_id} --path {os.path.join(output_dir,tar_dataset_path)}\")\n",
    "    print(upload_test_dataset_message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List the created datasets <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "message = subprocess.getoutput(f\"nvtl {model_name} list-datasets\")\n",
    "message = ast.literal_eval(message)\n",
    "for rsp in message:\n",
    "    rsp_keys = rsp.keys()\n",
    "    assert \"id\" in rsp_keys\n",
    "    assert \"type\" in rsp_keys\n",
    "    assert \"format\" in rsp_keys\n",
    "    assert \"name\" in rsp_keys\n",
    "    print(rsp[\"id\"],\"\\t\",rsp[\"type\"],\"\\t\",rsp[\"format\"],\"\\t\\t\",rsp[\"name\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an experiment <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "network_arch = model_name\n",
    "if \"classification_\" in network_arch:\n",
    "    encode_key = \"nvidia_tlt\"\n",
    "else:\n",
    "    encode_key = \"tlt_encode\"\n",
    "experiment_id = subprocess.getoutput(f\"nvtl {model_name} experiment-create --network_arch {network_arch} --encryption_key {encode_key} \")\n",
    "print(experiment_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assign train, eval datasets <a class=\"anchor\" id=\"head-10\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dataset_information = {\"train_datasets\":[train_dataset_id],\n",
    "                       \"eval_dataset\":eval_dataset_id,\n",
    "                       \"inference_dataset\":test_dataset_id,\n",
    "                       \"calibration_dataset\":train_dataset_id}\n",
    "patched_model = subprocess.getoutput(f\"nvtl {model_name} patch-artifact-metadata --id {experiment_id} --job_type experiment --update_info '{json.dumps(dataset_information)}' \")\n",
    "print(patched_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List experiments <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# List all pretrained models for the chosen network architecture\n",
    "filter_params = {\"network_arch\": network_arch}\n",
    "message = subprocess.getoutput(f\"nvtl {model_name} list-experiments --filter_params '{json.dumps(filter_params)}'\")\n",
    "message = ast.literal_eval(message)\n",
    "for rsp in message:\n",
    "    rsp_keys = rsp.keys()\n",
    "    if \"encryption_key\" not in rsp.keys():\n",
    "        assert \"name\" in rsp_keys and \"version\" in rsp_keys and \"ngc_path\" in rsp_keys and \"additional_id_info\" in rsp_keys\n",
    "        print(f'PTM Name: {rsp[\"name\"]}; PTM version: {rsp[\"version\"]}; NGC PATH: {rsp[\"ngc_path\"]}; Additional info: {rsp[\"additional_id_info\"]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assign PTM <a class=\"anchor\" id=\"head-7\"></a>\n",
    "\n",
    "Search for PTM on NGC for the Classification model chosen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Assigning pretrained models to different classification models\n",
    "# From the output of previous cell make the appropriate changes to this map if you want to change the default PTM backbone.\n",
    "# Changing the default backbone here requires changing default spec/config during train/eval etc like for example\n",
    "# If you are changing the ptm to resnet34, then you have to modify the config key num_layers if it exists to 34 manually\n",
    "pretrained_map = {\"classification_tf1\" : \"pretrained_classification:resnet18\",\n",
    "                  \"classification_tf2\" : \"pretrained_classification_tf2:efficientnet_b0\",\n",
    "                  \"classification_pyt\" : \"pretrained_fan_classification_imagenet:fan_hybrid_tiny\",\n",
    "                  \"multitask_classification\" : \"pretrained_classification:resnet10\"}\n",
    "no_ptm_models = set([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "if network_arch not in no_ptm_models:\n",
    "    filter_params = {\"network_arch\": network_arch}\n",
    "    message = subprocess.getoutput(f\"nvtl {model_name} list-experiments --filter_params '{json.dumps(filter_params)}'\")\n",
    "    message = ast.literal_eval(message)\n",
    "    ptm = []\n",
    "    for rsp in message:\n",
    "        rsp_keys = rsp.keys()\n",
    "        assert \"ngc_path\" in rsp_keys\n",
    "        if rsp[\"ngc_path\"].endswith(pretrained_map[network_arch]):\n",
    "            assert \"id\" in rsp_keys\n",
    "            ptm_id = rsp[\"id\"]\n",
    "            ptm = [ptm_id]\n",
    "            print(\"Metadata for model with requested NGC Path\")\n",
    "            print(rsp)\n",
    "            break\n",
    "    print(ptm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "if network_arch not in no_ptm_models:\n",
    "    ptm_information = {\"base_experiment\":ptm}\n",
    "    patched_model = subprocess.getoutput(f\"nvtl {model_name} patch-artifact-metadata --id {experiment_id} --job_type experiment --update_info '{json.dumps(ptm_information)}' \")\n",
    "    print(patched_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### View hyperparameters that are enabled for AutoML by default <a class=\"anchor\" id=\"head-11\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "if automl_enabled:\n",
    "    # View default automl specs enabled\n",
    "    ! nvtl {model_name} model-automl-defaults --id {experiment_id}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train <a class=\"anchor\" id=\"head-11\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "#### Set AutoML related configurations <a class=\"anchor\" id=\"head-12\"></a>\n",
    "Refer to these hyper-links to see the parameters supported by each network and add more parameters if necessary in addition to the default automl enabled parameters: \n",
    "\n",
    "[Classification TF1](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/classification_tf1/classification_tf1%20-%20train.csv), \n",
    "[Classification TF2](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/classification_tf2/classification_tf2%20-%20train.csv), \n",
    "[Classification Pytorch](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/classification_pyt/classification_pyt%20-%20train.csv), \n",
    "[Multitask classification](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/multitask_classification/multitask_classification%20-%20train.csv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "if automl_enabled:\n",
    "    # Choose any metric that is present in the kpi dictionary present in the model's status.json. \n",
    "    # Example status.json for each model can be found in the respective section in NVIDIA TAO DOCS here: https://docs.nvidia.com/tao/tao-toolkit/text/model_zoo/cv_models/index.html\n",
    "    if model_name == \"classification_pyt\":\n",
    "        metric = \"loss\"\n",
    "    else:\n",
    "        metric = \"kpi\" \n",
    "\n",
    "    additional_automl_parameters = [] #Refer to parameter list mentioned in the above links and add any extra parameter in addition to the default enabled ones\n",
    "    remove_default_automl_parameters = [] #Remove any hyperparameters that are enabled by default for AutoML\n",
    "\n",
    "    automl_information = {\"automl_enabled\":automl_enabled,\n",
    "                          \"automl_algorithm\":automl_algorithm,\n",
    "                          \"automl_max_recommendations\": 20, # Only for bayesian\n",
    "                          \"automl_R\": 27, # Only for hyperband\n",
    "                          \"automl_nu\": 3, # Only for hyperband\n",
    "                          \"epoch_multiplier\": 1, # Only for hyperband\n",
    "                          # Enable this if you want to add parameters to automl_add_hyperparameters below that are disabled by TAO in the automl_enabled column of the spec csv.\n",
    "                          # Warning: The parameters that are disabled are not tested by TAO, so there might be unexpected behaviour in overriding this\n",
    "                          \"override_automl_disabled_params\": False,\n",
    "                          \"metric\":metric,\n",
    "                          \"automl_add_hyperparameters\":str(additional_automl_parameters),\n",
    "                          \"automl_remove_hyperparameters\":str(remove_default_automl_parameters)\n",
    "                         }\n",
    "    patched_model = subprocess.getoutput(f\"nvtl {model_name} patch-artifact-metadata --id {experiment_id} --job_type experiment --update_info '{json.dumps(automl_information)}' \")\n",
    "    patched_model = json.loads(patched_model)\n",
    "    print(json.dumps(patched_model, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "#### Provide train specs <a class=\"anchor\" id=\"head-13\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Default train model specs\n",
    "train_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --action train --job_type experiment --id {experiment_id}\")\n",
    "train_specs = json.loads(train_specs)\n",
    "print(json.dumps(train_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Customize train model specs\n",
    "# Example for multitask_classification (for each network the parameter key might be different)\n",
    "if model_name == \"multitask_classification\":\n",
    "    train_specs[\"training_config\"][\"num_epochs\"] = 10\n",
    "    train_specs[\"gpus\"] = 1\n",
    "# Example for classification_pyt\n",
    "elif model_name == \"classification_pyt\":\n",
    "    train_specs[\"train\"][\"train_config\"][\"runner\"][\"max_epochs\"] = 40\n",
    "    train_specs[\"train\"][\"num_gpus\"] = 1\n",
    "    train_specs[\"gpus\"] = 1\n",
    "# Example for classification_tf1\n",
    "elif model_name == \"classification_tf1\":\n",
    "    train_specs[\"train_config\"][\"n_epochs\"] = 80\n",
    "    train_specs[\"gpus\"] = 1\n",
    "# Example for classification_tf2\n",
    "elif model_name == \"classification_tf2\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 80\n",
    "    train_specs[\"gpus\"] = 1\n",
    "print(json.dumps(train_specs, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "#### Run train <a class=\"anchor\" id=\"head-14\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action train --id {experiment_id} --specs '{json.dumps(train_specs)}'\")\n",
    "job_map[\"train_\" + model_name] = job_id\n",
    "print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Monitor job status\n",
    "if automl_enabled:    \n",
    "    while True:\n",
    "        clear_output(wait=True)\n",
    "        response = subprocess.getoutput(f\"nvtl {model_name} get-action-status --job_type experiment --id {experiment_id} --job {job_id}\")\n",
    "        response = json.loads(response)\n",
    "        if \"error_desc\" in response.keys() and response[\"error_desc\"] in (\"Job trying to retrieve not found\", \"No AutoML run found\"):\n",
    "            print(\"Job is being created\")\n",
    "            time.sleep(5)\n",
    "            continue\n",
    "        print(json.dumps(response, sort_keys=True, indent=4))\n",
    "        assert \"status\" in response.keys() and response.get(\"status\") != \"Error\"\n",
    "        if response.get(\"status\") in [\"Done\",\"Error\"]:\n",
    "            break\n",
    "        time.sleep(15)\n",
    "else:\n",
    "    # Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "    status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## To Stop an AutoML JOB\n",
    "#    1. Stop the 'Monitor job status' cell (the cell right before this cell) manually\n",
    "#    2. Uncomment the snippet in the next cell and run the cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# if automl_enabled:\n",
    "#     job_id = job_map[\"train_\" + model_name]\n",
    "#     job_id = subprocess.getoutput(f\"nvtl {model_name} job-cancel --job_type experiment --id {experiment_id} --job {job_id}\")\n",
    "#     job_map[\"canceled_\" + model_name] = job_id\n",
    "#     print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Resume AutoML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Uncomment the below snippet if you want to resume an already stopped AutoML job and then run the 'Monitor job status' cell above (4th cell above from this cell)\n",
    "# if automl_enabled:\n",
    "#     job_id = job_map[\"train_\" + model_name]\n",
    "#     job_id = subprocess.getoutput(f\"nvtl {model_name} job-resume --id {experiment_id} --job {job_id} --specs '{json.dumps(train_specs)}'\")\n",
    "#     job_map[\"resumed_\" + model_name] = job_id\n",
    "#     print(job_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Download train job artifacts <a class=\"anchor\" id=\"head-15\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "job_id = job_map[\"train_\" + model_name]\n",
    "file_list = subprocess.getoutput(f\"nvtl {model_name} list-job-files --id {experiment_id} --job {job_id} --job_type experiment --retrieve_logs True --retrieve_specs False\")\n",
    "print(file_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Patch the model with proper metric before training to run this cell; By default loss is used, but some models dont log the parameter under the name 'loss'\n",
    "# file_lists = []\n",
    "# temptar = subprocess.getoutput(f\"nvtl {model_name} download-selective-files --id {experiment_id} --job {job_id} --job_type experiment --workdir {workdir} --file_lists '{file_lists}' --best_model False --latest_model True --tar_files True\")\n",
    "# tar_command = f'tar -xvf {temptar} -C {workdir}/'\n",
    "# os.system(tar_command)\n",
    "# os.remove(temptar)\n",
    "# print(f\"Results at {workdir}/{job_id}\")\n",
    "# model_downloaded_path = f\"{workdir}/{job_id}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Downloading train job takes a longer time, uncomment this cell if you want to still proceed\n",
    "if download_jobs:\n",
    "    temptar = subprocess.getoutput(f\"nvtl {model_name} download-entire-job --id {experiment_id} --job {job_id} --job_type experiment --workdir {workdir}\")\n",
    "    tar_command = f'tar -xvf {temptar} -C {workdir}/'\n",
    "    os.system(tar_command)\n",
    "    os.remove(temptar)\n",
    "    print(f\"Results at {workdir}/{job_id}\")\n",
    "    model_downloaded_path = f\"{workdir}/{job_id}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# View the checkpoints generated for the training job and for automl jobs, in addition view: best performing model's config and the results of all automl experiments\n",
    "if download_jobs:\n",
    "    if automl_enabled:\n",
    "        !python3 -m pip install pandas==1.5.1\n",
    "        import pandas as pd\n",
    "        model_downloaded_path = f\"{model_downloaded_path}/best_model\"\n",
    "        assert glob.glob(f\"{model_downloaded_path}/*.protobuf\") or glob.glob(f\"{model_downloaded_path}/*.yaml\")\n",
    "\n",
    "    assert os.path.exists(model_downloaded_path)\n",
    "    assert (glob.glob(model_downloaded_path + \"/**/*.tlt\", recursive=True) + glob.glob(model_downloaded_path + \"/**/*.hdf5\", recursive=True) + glob.glob(model_downloaded_path + \"/**/*.pth\", recursive=True))\n",
    "\n",
    "    if os.path.exists(model_downloaded_path):        \n",
    "        #List the binary model file\n",
    "        print(\"\\nCheckpoints for the training experiment\")\n",
    "        if os.path.exists(model_downloaded_path+\"/train/weights\") and len(os.listdir(model_downloaded_path+\"/train/weights\")) > 0:\n",
    "            print(f\"Folder: {model_downloaded_path}/train/weights\")\n",
    "            print(\"Files:\", os.listdir(model_downloaded_path+\"/train/weights\"))\n",
    "        elif os.path.exists(model_downloaded_path+\"/weights\") and len(os.listdir(model_downloaded_path+\"/weights\")) > 0:\n",
    "            print(f\"Folder: {model_downloaded_path}/weights\")\n",
    "            print(\"Files:\", os.listdir(model_downloaded_path+\"/weights\"))\n",
    "        else:\n",
    "            print(f\"Folder: {model_downloaded_path}\")\n",
    "            print(\"Files:\", os.listdir(model_downloaded_path))\n",
    "\n",
    "        if automl_enabled:\n",
    "            assert glob.glob(f\"{model_downloaded_path}/*.protobuf\") or glob.glob(f\"{model_downloaded_path}/*.yaml\")\n",
    "            experiment_artifacts = json.load(open(f\"{model_downloaded_path}/controller.json\",\"r\"))\n",
    "            data_frame = pd.DataFrame(experiment_artifacts)\n",
    "            # Print experiment id/number and the corresponding result\n",
    "            print(\"\\nResults of all experiments\")\n",
    "            with pd.option_context('display.max_rows', None, 'display.max_columns', None, 'display.max_colwidth', None):\n",
    "                print(data_frame[[\"id\",\"result\"]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate <a class=\"anchor\" id=\"head-16\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Provide evaluate specs <a class=\"anchor\" id=\"head-16\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default evaluate model specs\n",
    "eval_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --action evaluate --job_type experiment --id {experiment_id}\")\n",
    "eval_specs = json.loads(eval_specs)\n",
    "print(json.dumps(eval_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize evaluate model specs\n",
    "# Change any spec if you wish\n",
    "print(json.dumps(eval_specs, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run evaluate <a class=\"anchor\" id=\"head-17\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Print model handler parameters\n",
    "model_parameters = subprocess.getoutput(f\"nvtl {model_name} get-metadata --id {experiment_id} --job_type experiment\")\n",
    "model_parameters = json.loads(model_parameters)\n",
    "update_checkpoint_choosing = {}\n",
    "update_checkpoint_choosing[\"checkpoint_choose_method\"] = model_parameters[\"checkpoint_choose_method\"]\n",
    "update_checkpoint_choosing[\"checkpoint_epoch_number\"] = model_parameters[\"checkpoint_epoch_number\"]\n",
    "print(json.dumps(update_checkpoint_choosing, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Change the method by which checkpoint from the parent action is chosen, when parent action is a train/retrain action.\n",
    "# Example for evaluate action below, can be applied in the same way for other actions too\n",
    "update_checkpoint_choosing[\"checkpoint_choose_method\"] = \"latest_model\" # Choose between best_model/latest_model/from_epoch_number\n",
    "# If from_epoch_number is chosen then assign the epoch number to the dictionary key in the format 'from_epoch_number{train_job_id}'\n",
    "# update_checkpoint_choosing[\"checkpoint_epoch_number\"][\"from_epoch_number_c2f76eb7-2a75-4197-9a84-c1547f20c17d\"] = 6\n",
    "\n",
    "patched_model = subprocess.getoutput(f\"nvtl {model_name} patch-artifact-metadata --id {experiment_id} --job_type experiment --update_info '{json.dumps(update_checkpoint_choosing)}'\")\n",
    "patched_model = json.loads(patched_model)\n",
    "print(json.dumps(patched_model, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "parent = job_map[\"train_\" + model_name]\n",
    "job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action evaluate --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(eval_specs)}'\")\n",
    "job_map[\"eval_\" + model_name] = job_id\n",
    "print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prune, Retrain and Evaluation <a class=\"anchor\" id=\"head-13\"></a>\n",
    "\n",
    "- We optimize the trained model by pruning and retraining in the following cells"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "#### Prune <a class=\"anchor\" id=\"head-18\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name != \"classification_pyt\":\n",
    "    # Default prune model specs\n",
    "    prune_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --id {experiment_id} --action prune --job_type experiment\")\n",
    "    prune_specs = json.loads(prune_specs)\n",
    "    print(json.dumps(prune_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "if model_name != \"classification_pyt\":\n",
    "    # Customize prune model specs\n",
    "    # Apply changes to specs dictionary if required here\n",
    "    if model_name == \"classification_tf2\":\n",
    "        prune_specs[\"prune\"][\"byom_model_path\"] = \"\"\n",
    "    print(json.dumps(prune_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if model_name != \"classification_pyt\":\n",
    "    parent = job_map[\"train_\" + model_name]\n",
    "    job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action prune --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(prune_specs)}'\")\n",
    "    job_map[\"prune_\" + model_name] = job_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check status of pruning job (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "if model_name != \"classification_pyt\":\n",
    "    status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "#### Retrain <a class=\"anchor\" id=\"head-20\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name != \"classification_pyt\":\n",
    "    # Default retrain model specs\n",
    "    retrain_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --id {experiment_id} --action retrain --job_type experiment\")\n",
    "    retrain_specs = json.loads(retrain_specs)\n",
    "    print(json.dumps(retrain_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name != \"classification_pyt\":\n",
    "    # Override any of the parameters listed in the previous cell as required\n",
    "    # Example for multitask_classification (for each network the parameter key might be different)\n",
    "    if model_name == \"multitask_classification\":\n",
    "        retrain_specs[\"training_config\"][\"num_epochs\"] = 10\n",
    "        retrain_specs[\"gpus\"] = 1\n",
    "    # Example for classification_tf1\n",
    "    elif model_name == \"classification_tf1\":\n",
    "        retrain_specs[\"train_config\"][\"n_epochs\"] = 80\n",
    "        retrain_specs[\"gpus\"] = 1\n",
    "    # Example for classification_tf2\n",
    "    elif model_name == \"classification_tf2\":\n",
    "        retrain_specs[\"train\"][\"num_epochs\"] = 80\n",
    "        retrain_specs[\"gpus\"] = 1\n",
    "    print(json.dumps(retrain_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if model_name != \"classification_pyt\":\n",
    "    parent = job_map[\"prune_\" + model_name]\n",
    "    job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action retrain --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(retrain_specs)}'\")\n",
    "    job_map[\"retrain_\" + model_name] = job_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check status of evaluate after retrain job(the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "if model_name != \"classification_pyt\":\n",
    "    status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Evaluate after retrain <a class=\"anchor\" id=\"head-20\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Default evaluate model specs\n",
    "if model_name != \"classification_pyt\":\n",
    "    eval_retrain_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --action evaluate --job_type experiment --id {experiment_id}\")\n",
    "    eval_retrain_specs = json.loads(eval_retrain_specs)\n",
    "    print(json.dumps(eval_retrain_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Customize evaluate model specs\n",
    "# Change any spec if you wish\n",
    "if model_name != \"classification_pyt\":\n",
    "    print(json.dumps(eval_retrain_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if model_name != \"classification_pyt\":\n",
    "    parent = job_map[\"retrain_\" + model_name]\n",
    "    job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action evaluate --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(eval_retrain_specs)}'\")\n",
    "    job_map[\"eval2_\" + model_name] = job_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check status of retrain job (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "if model_name != \"classification_pyt\":\n",
    "    status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Export <a class=\"anchor\" id=\"head-22\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Provide export specs <a class=\"anchor\" id=\"head-22\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default export model specs\n",
    "export_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --id {experiment_id} --action export --job_type experiment\")\n",
    "export_specs = json.loads(export_specs)\n",
    "print(json.dumps(export_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize export model specs\n",
    "# Apply changes to the specs dictionary here if required\n",
    "print(json.dumps(export_specs, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run export <a class=\"anchor\" id=\"head-23\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "parent = job_map[\"train_\" + model_name]\n",
    "job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action export --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(export_specs)}'\")\n",
    "job_map[\"export_\" + model_name] = job_id\n",
    "print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TRT Engine generation using TAO-Deploy <a class=\"anchor\" id=\"head-26\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Provide trt engine generation specs <a class=\"anchor\" id=\"head-26\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default gen_trt_engine model specs\n",
    "tao_deploy_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --id {experiment_id} --action gen_trt_engine --job_type experiment\")\n",
    "tao_deploy_specs = json.loads(tao_deploy_specs)\n",
    "print(json.dumps(tao_deploy_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize gen_trt_engine model specs\n",
    "if model_name == \"classification_tf2\":\n",
    "    tao_deploy_specs[\"gen_trt_engine\"][\"tensorrt\"][\"data_type\"] = \"int8\"\n",
    "elif model_name == \"classification_pyt\":\n",
    "    tao_deploy_specs[\"gen_trt_engine\"][\"tensorrt\"][\"data_type\"] = \"fp16\"\n",
    "else:\n",
    "    tao_deploy_specs[\"data_type\"] = \"int8\"\n",
    "print(json.dumps(tao_deploy_specs, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run TRT Engine generation using TAO-Deploy <a class=\"anchor\" id=\"head-27\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "parent = job_map[\"export_\" + model_name]\n",
    "job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action gen_trt_engine --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(tao_deploy_specs)}'\")\n",
    "job_map[\"gen_trt_engine_\" + model_name] = job_id\n",
    "print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TAO inference <a class=\"anchor\" id=\"head-28\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Provide TAO inference specs <a class=\"anchor\" id=\"head-28\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default inference model specs\n",
    "tao_inference_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --id {experiment_id} --action inference --job_type experiment\")\n",
    "tao_inference_specs = json.loads(tao_inference_specs)\n",
    "print(json.dumps(tao_inference_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize TAO inference specs\n",
    "#Apply changes to the specs dictionary here if required\n",
    "print(json.dumps(tao_inference_specs, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run TAO inference <a class=\"anchor\" id=\"head-29\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "parent = job_map[\"train_\" + model_name]\n",
    "job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action inference --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(tao_inference_specs)}'\")\n",
    "job_map[\"tlt_inference_\" + model_name] = job_id\n",
    "print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if download_jobs:\n",
    "    temptar = subprocess.getoutput(f\"nvtl {model_name} download-entire-job --id {experiment_id} --job {job_id} --job_type experiment --workdir {workdir}\")\n",
    "    tar_command = f'tar -xvf {temptar} -C {workdir}/'\n",
    "    os.system(tar_command)\n",
    "    os.remove(temptar)\n",
    "    print(f\"Results at {workdir}/{job_id}\")\n",
    "    inference_out_path = f\"{workdir}/{job_id}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if download_jobs:\n",
    "    # Print Classification results\n",
    "    if model_name == \"classification_tf1\":\n",
    "        assert os.path.exists(f'{inference_out_path}/result.csv')\n",
    "        !cat {inference_out_path}/result.csv\n",
    "    elif \"classification_\" in model_name:\n",
    "        assert os.path.exists(f'{inference_out_path}/inference/result.csv')\n",
    "        !cat {inference_out_path}/inference/result.csv\n",
    "    elif model_name == \"multitask_classification\":\n",
    "        assert os.path.exists(f'{inference_out_path}/result.txt')\n",
    "        !cat {inference_out_path}/result.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TRT inference <a class=\"anchor\" id=\"head-30\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Provide TRT inference specs <a class=\"anchor\" id=\"head-30\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default inference model specs\n",
    "trt_inference_specs = subprocess.getoutput(f\"nvtl {model_name} get-spec --id {experiment_id} --action inference --job_type experiment\")\n",
    "trt_inference_specs = json.loads(trt_inference_specs)\n",
    "print(json.dumps(trt_inference_specs, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize TAO inference specs\n",
    "# Apply changes to the specs dictionary here if required\n",
    "print(json.dumps(trt_inference_specs, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run TRT inference <a class=\"anchor\" id=\"head-31\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default inference model specs\n",
    "parent = job_map[\"gen_trt_engine_\" + model_name]\n",
    "job_id = subprocess.getoutput(f\"nvtl {model_name} experiment-run-action --action inference --id {experiment_id} --parent_job_id {parent} --specs '{json.dumps(trt_inference_specs)}'\")\n",
    "job_map[\"trt_inference_\" + model_name] = job_id\n",
    "print(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "status = my_tail(model_name, experiment_id, job_id, \"experiment\", workdir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if download_jobs:\n",
    "    temptar = subprocess.getoutput(f\"nvtl {model_name} download-entire-job --id {experiment_id} --job {job_id} --job_type experiment --workdir {workdir}\")\n",
    "    tar_command = f'tar -xvf {temptar} -C {workdir}/'\n",
    "    os.system(tar_command)\n",
    "    os.remove(temptar)\n",
    "    print(f\"Results at {workdir}/{job_id}\")\n",
    "    inference_out_path = f\"{workdir}/{job_id}\"\n",
    "    assert glob.glob(f\"{inference_out_path}/**/*result.csv\", recursive=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Print Classification results\n",
    "if download_jobs:\n",
    "    if model_name in (\"classification_tf1\", \"multitask_classification\"):\n",
    "        !cat {inference_out_path}/result.csv\n",
    "    elif \"classification_\" in model_name:\n",
    "        !cat {inference_out_path}/inference/result.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete model <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subprocess.getoutput(f\"nvtl {model_name} experiment-delete --id {experiment_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Delete train dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subprocess.getoutput(f\"nvtl {model_name} dataset-delete --id {train_dataset_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Delete val dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subprocess.getoutput(f\"nvtl {model_name} dataset-delete --id {eval_dataset_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Delete test dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subprocess.getoutput(f\"nvtl {model_name} dataset-delete --id {test_dataset_id}\")"
   ]
  }
 ],
 "metadata": {
  "datalore": {
   "base_environment": "default",
   "computation_mode": "JUPYTER",
   "package_manager": "pip",
   "packages": [],
   "version": 1
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
